package com.yifeng.repo.controller.rocketmq.message.consume;

import com.yifeng.repo.base.utils.common.BaseUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.consumer.rebalance.AbstractAllocateMessageQueueStrategy;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMachineRoomNearby;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueAveragely;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueByMachineRoom;
import org.apache.rocketmq.common.message.MessageQueue;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * Created by daibing on 2023/10/21.
 */
@Slf4j
public class AllocateMessageQueueByMsgType extends AbstractAllocateMessageQueueStrategy {
    private static final Long ONE_DAY_MILLIS = 24 * 60 * 60 * 1000L;

    private final ConcurrentMap<String, List<Integer>> cid2QueueIds = new ConcurrentHashMap<>();
    private final AtomicLong lastPrintTime = new AtomicLong(System.currentTimeMillis());

    private final String msgType;
    private final Map<Integer, String> queueId2MsgType;

    public AllocateMessageQueueByMsgType(String msgType, Map<Integer, String> queueId2MsgType) {
        this.msgType = msgType;
        this.queueId2MsgType = queueId2MsgType;
    }

    /**
     * 按照消息类型分配消息队列策略
     * 1、每个消费者仅消费一个类型的消息队列，避免各个类型的消息互相影响
     * 2、借鉴多种算法的思想
     *
     * @see AllocateMessageQueueByMachineRoom 机房编号分配
     * @see AllocateMessageQueueAveragely 平均分配
     * @see AllocateMachineRoomNearby 同机房优先分配
     */
    @Override
    public List<MessageQueue> allocate(String consumerGroup, String currentCID, List<MessageQueue> mqAll, List<String> cidAll) {
        // 1. 按照消息类型筛选，特别注意创建cid（消费者instanceName）指定名称规则和这里过滤名称规则要保持一致。
        List<MessageQueue> mqAllByType = mqAll.stream()
                .filter(q -> queueId2MsgType.containsKey(q.getQueueId()) && queueId2MsgType.get(q.getQueueId()).equals(msgType))
                .collect(Collectors.toList());
        List<String> cidAllByType = cidAll.stream()
                .filter(c -> c.contains("-" + msgType + "-" + consumerGroup))
                .collect(Collectors.toList());

        // 2. 检查参数，不合理的参数就不做分配了
        List<MessageQueue> result = new ArrayList<>();
        if (!check(consumerGroup, currentCID, mqAllByType, cidAllByType)) {
            return checkAndUpdateAllocateMessageQueue(currentCID, result);
        }

        // 3. 获取当前消费者ID在cidAll中的索引位置，以及计算取模结果
        int index = cidAllByType.indexOf(currentCID);
        int mod = mqAllByType.size() % cidAllByType.size();

        // 4. 计算每隔消费者平均分配到的数量，以及计算起始索引值
        int averageSize =
                mqAllByType.size() <= cidAllByType.size() ? 1 : (mod > 0 && index < mod ? mqAllByType.size() / cidAllByType.size()
                        + 1 : mqAllByType.size() / cidAllByType.size());
        int startIndex = (mod > 0 && index < mod) ? index * averageSize : index * averageSize + mod;

        // 5. 计算目标MessageQueue
        int range = Math.min(averageSize, mqAllByType.size() - startIndex);
        for (int i = 0; i < range; i++) {
            result.add(mqAllByType.get((startIndex + i) % mqAllByType.size()));
        }
        return checkAndUpdateAllocateMessageQueue(currentCID, result);
    }

    private List<MessageQueue> checkAndUpdateAllocateMessageQueue(String currentCID, List<MessageQueue> allocateMessageQueues) {
        // 获取当前配置和最新配置
        List<Integer> currentMessageQueueIds = cid2QueueIds.get(currentCID);
        List<Integer> newMessageQueueIds = allocateMessageQueues.stream().map(MessageQueue::getQueueId).collect(Collectors.toList());

        // 判断配置未发生变更，每天打印一次
        if (currentMessageQueueIds != null
                && currentMessageQueueIds.containsAll(newMessageQueueIds)
                && newMessageQueueIds.containsAll(currentMessageQueueIds)) {
            if (lastPrintTime.get() + ONE_DAY_MILLIS <= System.currentTimeMillis()) {
                lastPrintTime.set(System.currentTimeMillis());
                log.info("MQ分配打印: ownerId={}, ownerName={}, msgType={}, currentCID={}, currentMessageQueueIds={}",
                        BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME, msgType, currentCID, currentMessageQueueIds);
            }
            return allocateMessageQueues;
        }

        // 已经发送变更就更新缓存并且打印日志
        cid2QueueIds.put(currentCID, newMessageQueueIds);
        log.info("MQ分配变更: ownerId={}, ownerName={}, msgType={}, currentCID={}, currentMessageQueueIds={}, newMessageQueueIds={}",
                BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME, msgType, currentCID, currentMessageQueueIds, newMessageQueueIds);
        return allocateMessageQueues;
    }

    @Override
    public String getName() {
        return msgType;
    }
}
